import numpy as np
import scipy.linalg as LA


class TLSESPRIT:
    def __init__(self, num_signals, num_sensors, sensor_distance, wavelength):
        self.num_signals = num_signals
        self.num_sensors = num_sensors
        self.sensor_distance = sensor_distance
        self.wavelength = wavelength

    def awgn(self, x, snr):
        spower = np.sum((np.abs(x) ** 2)) / x.size
        x = x + np.sqrt(spower / snr) * (np.random.randn(x.shape[0], x.shape[1]) + 1j * np.random.randn(x.shape[0], x.shape[1]))
        return x

    def TLS_ESPRIT(self, array_signal):
        #计算协方差矩阵
        R = np.dot(array_signal, array_signal.conj().T) / array_signal.shape[1]
        # 特征值分解
        D, EV = LA.eigh(R)
        index = np.argsort(D)
        EN = EV[:, index][:, -self.num_signals:]

        Exy = np.concatenate((EN[:-1, :], EN[1:, :]), axis=1)
        E_xy = Exy.conj().T @ Exy
        D_xy, EV_xy = LA.eig(E_xy)
        index1 = np.argsort(D_xy)
        EV_xy = EV_xy[:, index1]
        E12 = EV_xy[:self.num_signals, :self.num_signals]
        E22 = EV_xy[self.num_signals:, :self.num_signals]
        Psi = -E12 @ LA.pinv(E22)

        D_si, V_si = LA.eig(Psi)
        ephi = np.angle(D_si)
        angles = -np.arcsin(ephi / (2 * np.pi) / (self.sensor_distance / self.wavelength)) * 180 / np.pi
        return np.sort(angles)

# 使用示例
num_signals = 3
num_sensors = 8
sensor_distance = 0.5
wavelength = 1

tls_esprit = TLSESPRIT(num_signals, num_sensors, sensor_distance, wavelength)

# 生成仿真信号
array_signal, true_angles = generate_signal(num_signals=num_signals, num_sensors=num_sensors)

# 添加噪声
SNR = 10
array_signal = tls_esprit.awgn(array_signal, SNR)

# 执行TLS-ESPRIT算法
estimated_angles = tls_esprit.TLS_ESPRIT(array_signal)
print("True Angles:", true_angles)
print("Estimated Angles:", estimated_angles)